/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.fnexecution.logging;



import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.LogControl;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnLoggingGrpc;

import com.bff.gaia.unified.runners.fnexecution.FnService;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.io.grpc.stub.StreamObserver;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableSet;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.util.Set;

import java.util.concurrent.ConcurrentHashMap;

import java.util.concurrent.ConcurrentMap;



/** An implementation of the Unified Fn Logging Service over gRPC. */

public class GrpcLoggingService extends UnifiedFnLoggingGrpc.UnifiedFnLoggingImplBase

    implements FnService {

  private static final Logger LOGGER = LoggerFactory.getLogger(GrpcLoggingService.class);



  public static GrpcLoggingService forWriter(LogWriter writer) {

    return new GrpcLoggingService(writer);

  }



  private final LogWriter logWriter;

  private final ConcurrentMap<InboundObserver, StreamObserver<LogControl>> connectedClients;



  private GrpcLoggingService(LogWriter logWriter) {

    this.logWriter = logWriter;

    connectedClients = new ConcurrentHashMap<>();

  }



  @Override

  public void close() throws Exception {

    Set<InboundObserver> remainingClients = ImmutableSet.copyOf(connectedClients.keySet());

    if (!remainingClients.isEmpty()) {

      LOGGER.info(

          "{} Unified Fn Logging clients still connected during shutdown.", remainingClients.size());



      // Signal server shutting down to all remaining connected clients.

      for (InboundObserver client : remainingClients) {

        // We remove these from the connected clients map to prevent a race between

        // this close method and the InboundObserver calling a terminal method on the

        // StreamObserver. If we removed it, then we are responsible for the terminal call.

        completeIfNotNull(connectedClients.remove(client));

      }

    }

  }



  @Override

  public StreamObserver<UnifiedFnApi.LogEntry.List> logging(

      StreamObserver<UnifiedFnApi.LogControl> outboundObserver) {

    LOGGER.info("Unified Fn Logging client connected.");

    InboundObserver inboundObserver = new InboundObserver();

    connectedClients.put(inboundObserver, outboundObserver);

    return inboundObserver;

  }



  private void completeIfNotNull(StreamObserver<UnifiedFnApi.LogControl> outboundObserver) {

    if (outboundObserver != null) {

      try {

        outboundObserver.onCompleted();

      } catch (RuntimeException ignored) {

        // Completing outbound observer failed, ignoring failure and continuing

        LOGGER.warn("Unified Fn Logging client failed to be complete.", ignored);

      }

    }

  }



  /**

   * An inbound {@link StreamObserver} that forwards incoming messages to the client logger.

   *

   * <p>Mutually hangs up on clients that have errored or completed.

   */

  private class InboundObserver implements StreamObserver<UnifiedFnApi.LogEntry.List> {

    @Override

    public void onNext(UnifiedFnApi.LogEntry.List value) {

      for (UnifiedFnApi.LogEntry logEntry : value.getLogEntriesList()) {

        logWriter.log(logEntry);

      }

    }



    @Override

    public void onError(Throwable t) {

      LOGGER.warn("Logging client failed unexpectedly.", t);

      // We remove these from the connected clients map to prevent a race between

      // the close method and this InboundObserver calling a terminal method on the

      // StreamObserver. If we removed it, then we are responsible for the terminal call.

      completeIfNotNull(connectedClients.remove(this));

    }



    @Override

    public void onCompleted() {

      LOGGER.info("Logging client hanged up.");

      // We remove these from the connected clients map to prevent a race between

      // the close method and this InboundObserver calling a terminal method on the

      // StreamObserver. If we removed it, then we are responsible for the terminal call.

      completeIfNotNull(connectedClients.remove(this));

    }

  }

}